import torch, torchaudio
from modelscope.pipelines import pipeline # type: ignore
from modelscope.utils.constant import Tasks # type: ignore

class DeNoise:
    def __init__(self, model_dir, logLevel='info'):
        self.deNoiseModel = pipeline(
            Tasks.acoustic_noise_suppression,
            model=f'{model_dir}',
            disable_update=True,
            disable_pbar=True if logLevel != 'debug' else False,
            disable_log=True if logLevel != 'debug' else False,
        )
        
    def deNoiseFunc(self, input, output = None, gain_factor=2, target_rate=16000):
        if output is None:
            output = input
        waveform, sample_rate = torchaudio.load(input)
        if (sample_rate != target_rate):
            # 创建Resample对象
            resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_rate)
            # 重采样
            resampled_waveform = resampler(waveform)
            torchaudio.save(input, resampled_waveform, target_rate)
        self.deNoiseModel(input, output_path=output)
        if gain_factor != 1:
            waveform, sample_rate = torchaudio.load(output)
            resampled_waveform = waveform * gain_factor
            resampled_waveform = torch.clamp(resampled_waveform, -1.0, 1.0)
            torchaudio.save(output, resampled_waveform, target_rate)